#include "des.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#define ONES(k) ((1ULL << (k)) - 1)

#define NUM_SUBKEYS 16
#define BITS_PC_1 56
#define BITS_PC_2 48
#define BITS_IP 64
#define BITS_E 48
#define BITS_P 32
#define BITS_FP 64

/* Tables */
static int table_initialized = 0;
static int PC_1[BITS_PC_1];
static int LEFT_SHIFTS[NUM_SUBKEYS];
static int PC_2[BITS_PC_2];
static int IP[BITS_IP];
static int E[BITS_E];
static int S[8][4][16];
static int P[BITS_P];
static int FP[BITS_FP];

static uint64_t K[NUM_SUBKEYS + 1];

void set_bit(uint64_t *num, size_t i) {
    *num |= ((uint64_t)1 << i);
}
int get_bit(uint64_t num, size_t i) {
    return (num >> i) & 1;
}

uint64_t rearrange(size_t bits_in, size_t bits_out, uint64_t in, const int *table) {
    uint64_t rt = 0;
    for (int i = 0; i < bits_out; ++i) {
        if (get_bit(in, bits_in - table[i])) {
            set_bit(&rt, bits_out - 1 - i);
        }
    }
    return rt;
}

void read_table(int *table, const char *filename, size_t size) {
    FILE *fp = fopen(filename, "r");
    for (int i = 0; i < size; ++i) {
        fscanf(fp, "%d", table + i);
    }
    fclose(fp);
}
void read_sboxes(int S[8][4][16]) {
    FILE *fp = fopen("tables/sboxes.txt", "r");
    for (int i = 0; i < 8; ++i) {
        for (int j = 0; j < 4; ++j) {
            for (int k = 0; k < 16; ++k) {
                fscanf(fp, "%d", S[i][j] + k);
            }
        }
    }
    fclose(fp);
}

void des_init(uint64_t key) {
    if (!table_initialized) {
        table_initialized = 1;
        read_table(PC_1, "tables/table_pc_1.txt", BITS_PC_1);
        read_table(LEFT_SHIFTS, "tables/left_shifts.txt", NUM_SUBKEYS);
        read_table(PC_2, "tables/table_pc_2.txt", BITS_PC_2);
        read_table(IP, "tables/table_ip.txt", BITS_IP);
        read_table(E, "tables/table_e.txt", BITS_E);
        read_sboxes(S);
        read_table(P, "tables/table_p.txt", BITS_P);
        read_table(FP, "tables/table_fp.txt", BITS_FP);
    }
    K[0] = rearrange(64, 56, key, PC_1);
    uint64_t *C = malloc(sizeof(uint64_t) * (NUM_SUBKEYS + 1));
    uint64_t *D = malloc(sizeof(uint64_t) * (NUM_SUBKEYS + 1));
    C[0] = ((K[0] >> 28) & ONES(28));
    D[0] = (K[0] & ONES(28));
    int sh;
    for (int i = 1; i <= NUM_SUBKEYS; ++i) {
        sh = LEFT_SHIFTS[i - 1];
        C[i] = ((C[i - 1] << sh) | ((C[i - 1] >> (28 - sh)) & ONES(sh))) & ONES(28);
        D[i] = ((D[i - 1] << sh) | ((D[i - 1] >> (28 - sh)) & ONES(sh))) & ONES(28);
        K[i] = rearrange(56, 48, (C[i] << 28) | D[i], PC_2);
    }
    free(C);
    free(D);
}

size_t get_number_of_encrypted_bytes(size_t size_message) {
    return (size_message / 8 + 1) * 8;
}

void des_encrypt_bytes(uint8_t *out, const uint8_t *in, size_t size) {
    size_t num_blocks = size / 8 + 1;
    uint64_t *in_blocks = (uint64_t*)malloc(sizeof(uint64_t) * num_blocks);
    uint64_t *out_blocks = (uint64_t*)malloc(sizeof(uint64_t) * num_blocks);
    memset(in_blocks, 0, sizeof(uint64_t) * num_blocks);
    for (int i = 0; i < size; ++i) {
        in_blocks[i / 8] |= ((uint64_t)in[i] << (8 * (7 - i % 8)));
    }
    size_t out_bytes = num_blocks * 8;
    size_t x = out_bytes - size;
    for (int i = 0; i < x; ++i) {
        in_blocks[num_blocks - 1] |= (x << (8 * i));
    }

    des_encrypt_blocks(out_blocks, in_blocks, num_blocks);

    for (int i = 0; i < out_bytes; ++i) {
        out[i] = (out_blocks[i / 8] >> (8 * (7 - i % 8))) & ONES(8);
    }
    free(in_blocks);
    free(out_blocks);
}

void des_encrypt_blocks(uint64_t *out, const uint64_t *in, size_t size) {
    for (size_t i = 0; i < size; ++i) {
        out[i] = des_encrypt_block(in[i]);
    }
}

uint64_t feistel(uint64_t r, uint64_t k) {
    uint64_t er = rearrange(32, 48, r, E);
    uint64_t ker = k ^ er;
    uint64_t rt = 0, x;
    int row, col;
    for (int i = 0; i < 8; ++i) {
        x = ((ker >> (6 * i)) & ONES(6));
        row = ((x >> 4) & 2) | (x & 1);
        col = ((x >> 1) & ONES(4));
        rt |= (S[7 - i][row][col] << (4 * i));
    }
    rt = rearrange(32, 32, rt, P);
    return rt;
}

uint64_t des_encrypt_block(uint64_t in) {
    uint64_t ip = rearrange(64, 64, in, IP);
    static uint64_t L[NUM_SUBKEYS + 1], R[NUM_SUBKEYS + 1];
    L[0] = ((ip >> 32) & ONES(32));
    R[0] = (ip & ONES(32));
    for (int i = 1; i <= NUM_SUBKEYS; ++i) {
        L[i] = R[i - 1];
        R[i] = L[i - 1] ^ feistel(R[i - 1], K[i]);
    }
    uint64_t rt = rearrange(64, 64, (R[NUM_SUBKEYS] << 32) | L[NUM_SUBKEYS], FP);
    return rt;
}

void des_decrypt_bytes(uint8_t *out, size_t *size_out, const uint8_t *in, size_t size_in) {
    size_t num_blocks = size_in / 8;
    uint64_t *in_blocks = (uint64_t*)malloc(sizeof(uint64_t) * num_blocks);
    uint64_t *out_blocks = (uint64_t*)malloc(sizeof(uint64_t) * num_blocks);
    memset(in_blocks, 0, sizeof(uint64_t) * num_blocks);
    for (int i = 0; i < size_in; ++i) {
        in_blocks[i / 8] |= ((uint64_t)in[i] << (8 * (7 - i % 8)));
    }

    des_decrypt_blocks(out_blocks, in_blocks, num_blocks);

    uint8_t x = out_blocks[num_blocks - 1] & ONES(8);
    *size_out = num_blocks * 8 - (size_t)x;
    for (int i = 0; i < *size_out; ++i) {
        out[i] = (out_blocks[i / 8] >> (8 * (7 - i % 8))) & ONES(8);
    }
    free(in_blocks);
    free(out_blocks);
}

void des_decrypt_blocks(uint64_t *out, const uint64_t *in, size_t size) {
    for (size_t i = 0; i < size; ++i) {
        out[i] = des_decrypt_block(in[i]);
    }
}

uint64_t des_decrypt_block(uint64_t in) {
    uint64_t ip = rearrange(64, 64, in, IP);
    static uint64_t L[NUM_SUBKEYS + 1], R[NUM_SUBKEYS + 1];
    L[NUM_SUBKEYS] = ((ip >> 32) & ONES(32));
    R[NUM_SUBKEYS] = (ip & ONES(32));
    for (int i = NUM_SUBKEYS - 1; i >= 0; --i) {
        L[i] = R[i + 1];
        R[i] = L[i + 1] ^ feistel(R[i + 1], K[i + 1]);
    }
    uint64_t rt = rearrange(64, 64, (R[0] << 32) | L[0], FP);
    return rt;
}
